# from openloong.alg import FootPlacement

import mujoco
import glfw
import numpy as np
from scipy.spatial.transform import Rotation

# from .. import openloong
from openloong.sim_interface import UIctr
from openloong.sim_interface import MJ_Interface
from openloong.common import PVT_Ctr
from openloong.common import DataLogger, DataBus, MotionState, LegState
from openloong.alg import Pin_KinDyn
from openloong.mathq.useful_math import eul2Rot, eigen2std
from openloong.alg import WBC_priority
from openloong.alg import MPC
from openloong.alg import nx, nu
from openloong.alg import GaitScheduler
from openloong.alg import FootPlacement
from openloong.alg import JoyStickInterpreter

model_path = "robot_models/scene_board.xml"
mj_model = mujoco.MjModel.from_xml_path(model_path)
mj_data = mujoco.MjData(mj_model)

def main():
    # // initialize classes
    # uiController = UIctr(mj_model,mj_data)  ## UI control for Mujoco
    # mj_interface = MJ_Interface(mj_model, mj_data) ## data interface for Mujoco
    # kinDynSolver = Pin_KinDyn("robot_models/AzureLoong.urdf")   ## kinematics and dynamics solver
    # RobotState = DataBus(kinDynSolver.model_nv)  ## data bus
    # WBC_solv = WBC_priority(kinDynSolver.model_nv, 18, 22, 0.7, mj_model.opt.timestep)  ## WBC solver
    # # MPC_solv = MPC(dt_200Hz)   ## mpc controller
    # gaitScheduler = GaitScheduler(0.25, mj_model.opt.timestep)   ##  gait scheduler
    # pvtCtr = PVT_Ctr(mj_model.opt.timestep, "user_examples/openloong/common/joint_ctrl_config.json")  ## PVT joint control
    # footPlacement = FootPlacement()  ## foot-placement planner
    # jsInterp = JoyStickInterpreter(mj_model.opt.timestep)  ##desired baselink velocity generator
    # logger = DataLogger("user_examples/record/datalog.log")  ## data logger
    
    uiController = UIctr(mj_model,mj_data)  ## UI control for Mujoco
    mj_interface = MJ_Interface(mj_model, mj_data) ## data interface for Mujoco
    kinDynSolver = Pin_KinDyn("robot_models/AzureLoong.urdf")   ## kinematics and dynamics solver
    RobotState = DataBus(kinDynSolver.model_nv)  ## data bus
    WBC_solv = WBC_priority(kinDynSolver.model_nv, 18, 22, 0.7, mj_model.opt.timestep)  ## WBC solver
    # MPC_solv = MPC(dt_200Hz)   ## mpc controller
    gaitScheduler = GaitScheduler(0.25, mj_model.opt.timestep)   ##  gait scheduler
    pvtCtr = PVT_Ctr(mj_model.opt.timestep, "user_examples/openloong/common/joint_ctrl_config.json")  ## PVT joint control
    footPlacement = FootPlacement()  ## foot-placement planner
    jsInterp = JoyStickInterpreter(mj_model.opt.timestep)  ##desired baselink velocity generator
    logger = DataLogger("user_examples/record/datalog.log")  ## data logger

    # Variables initialization
    stand_leg_length = 1.01  # desired baselink height
    foot_height = 0.07  # distance between the foot ankle joint and the bottom
    xv_des = 0.7  # desired velocity in x direction

    RobotState.width_hips = 0.229
    footPlacement.kp_vx = 0.03
    footPlacement.kp_vy = 0.035
    footPlacement.kp_wz = 0.03
    footPlacement.step_height = 0.25
    footPlacement.leg_length = stand_leg_length

    model_nv = kinDynSolver.model_nv

    # Initial position and posture for foot-end and hand
    motors_pos_des = np.zeros(model_nv - 6)
    motors_pos_cur = np.zeros(model_nv - 6)
    motors_vel_des = np.zeros(model_nv - 6)
    motors_vel_cur = np.zeros(model_nv - 6)
    motors_tau_des = np.zeros(model_nv - 6)
    motors_tau_cur = np.zeros(model_nv - 6)

    fe_l_pos_L_des = np.array([-0.018, 0.113, -stand_leg_length])
    fe_r_pos_L_des = np.array([-0.018, -0.116, -stand_leg_length])
    fe_l_eul_L_des = np.array([-0.000, -0.008, -0.000])
    fe_r_eul_L_des = np.array([0.000, -0.008, 0.000])
    fe_l_rot_des = eul2Rot(fe_l_eul_L_des[0], fe_l_eul_L_des[1], fe_l_eul_L_des[2])
    fe_r_rot_des = eul2Rot(fe_r_eul_L_des[0], fe_r_eul_L_des[1], fe_r_eul_L_des[2])

    hd_l_pos_L_des = np.array([-0.02, 0.32, -0.159])
    hd_r_pos_L_des = np.array([-0.02, -0.32, -0.159])
    hd_l_eul_L_des = np.array([-1.7581, 0.2129, 2.9581])
    hd_r_eul_L_des = np.array([1.7581, 0.2129, -2.9581])
    hd_l_rot_des = eul2Rot(hd_l_eul_L_des[0], hd_l_eul_L_des[1], hd_l_eul_L_des[2])
    hd_r_rot_des = eul2Rot(hd_r_eul_L_des[0], hd_r_eul_L_des[1], hd_r_eul_L_des[2])

    res_leg = kinDynSolver.computeInK_Leg(fe_l_rot_des, fe_l_pos_L_des, fe_r_rot_des, fe_r_pos_L_des)
    res_hand = kinDynSolver.computeInK_Hand(hd_l_rot_des, hd_l_pos_L_des, hd_r_rot_des, hd_r_pos_L_des)
    qIniDes = np.zeros(mj_model.nq)
    qIniDes[7:] = res_leg.jointPosRes + res_hand.jointPosRes
    WBC_solv.set_qini(qIniDes, RobotState.q)

    # Register variable names for data logger
    logger.add_item("simTime", 1)
    logger.add_item("motors_pos_cur", model_nv - 6)
    logger.add_item("motors_vel_cur", model_nv - 6)
    logger.add_item("rpy", 3)
    logger.add_item("fL", 3)
    logger.add_item("fR", 3)
    logger.add_item("basePos", 3)
    logger.add_item("baseLinVel", 3)
    logger.add_item("baseAcc", 3)
    logger.add_item("baseAngVel", 3)
    logger.finish_item_adding()

    # Simulation loop
    sim_end_time = 30
    sim_start = mj_data.time
    sim_time = mj_data.time
    start_stepping_time = 3
    start_walking_time = 5
    # start_stepping_time = 1
    # start_walking_time = 2

    # Initialize UI: GLFW
    uiController.iniGLFW()
    uiController.enableTracking()  # enable viewpoint tracking of the body 1 of the robot
    uiController.createWindow("Demo", False)

    # viewer_handle = Handle(cam=camera,  opt=mujoco.MjvOption(), user_scn=scene)

    while not glfw.window_should_close(uiController.window):
        sim_start = mj_data.time
        while (mj_data.time - sim_start < 1.0 / 60.0)  and uiController.runSim:  ## and viewer_handle.is_running():
            # mj_sim = MjSim(mj_model, mj_data)
            mujoco.mj_step(mj_model, mj_data)
            
            sim_time = mj_data.time
            print(f"-------------{sim_time:.3f} s-------------")

            mj_interface.updateSensorValues()
            mj_interface.dataBusWrite(RobotState)

            kinDynSolver.dataBusRead(RobotState)
            kinDynSolver.computeJ_dJ()
            kinDynSolver.computeDyn()
            kinDynSolver.dataBusWrite(RobotState)

            if sim_time > start_walking_time:
                jsInterp.setWzDesLPara(0, 1)
                jsInterp.setVxDesLPara(xv_des, 2.0)
                RobotState.motionState = MotionState.Walk  # start walking
            else:
                jsInterp.setIniPos(RobotState.q[0], RobotState.q[1], RobotState.base_rpy[2])

            jsInterp.step()
            RobotState.js_pos_des[2] = stand_leg_length + foot_height
            jsInterp.dataBusWrite(RobotState)

            if sim_time >= start_stepping_time:
                gaitScheduler.dataBusRead(RobotState)
                gaitScheduler.step()
                gaitScheduler.dataBusWrite(RobotState)

                footPlacement.dataBusRead(RobotState)
                footPlacement.getSwingPos()
                footPlacement.dataBusWrite(RobotState)

            # WBC input
            RobotState.Fr_ff = np.zeros(12)
            RobotState.des_ddq = np.zeros(mj_model.nv)
            RobotState.des_dq = np.zeros(mj_model.nv)
            RobotState.des_delta_q = np.zeros(mj_model.nv)
            RobotState.base_rpy_des = np.array([0, 0, jsInterp.thetaZ])
            RobotState.base_pos_des[2] = stand_leg_length + foot_height

            RobotState.Fr_ff = np.array([0, 0, 370, 0, 0, 0, 0, 0, 370, 0, 0, 0])

            if sim_time > start_walking_time + 1:
                RobotState.des_delta_q[:2] = jsInterp.vx_W * mj_model.opt.timestep, jsInterp.vy_W * mj_model.opt.timestep
                RobotState.des_delta_q[5] = jsInterp.wz_L * mj_model.opt.timestep
                RobotState.des_dq[:2] = jsInterp.vx_W, jsInterp.vy_W
                RobotState.des_dq[5] = jsInterp.wz_L

                k = 5.0
                RobotState.des_ddq[:2] = k * (jsInterp.vx_W - RobotState.dq[0]), k * (jsInterp.vy_W - RobotState.dq[0])
                RobotState.des_ddq[5] = k * (jsInterp.wz_L - RobotState.dq[5])

                # // WBC Calculation
            WBC_solv.dataBusRead(RobotState)
            WBC_solv.computeDdq(kinDynSolver)
            WBC_solv.computeTau()
            WBC_solv.dataBusWrite(RobotState)

            # // get the final joint command
            if (sim_time <= start_walking_time):
                RobotState.motors_pos_des= np.array(res_leg.jointPosRes + res_hand.jointPosRes)
                RobotState.motors_vel_des=motors_vel_des
                RobotState.motors_tor_des=motors_tau_des
            else:
                pos_des = kinDynSolver.integrateDIY(RobotState.q, RobotState.wbc_delta_q_final)
                RobotState.motors_pos_des = pos_des[7:].flatten()
                RobotState.motors_vel_des = RobotState.wbc_dq_final.flatten()
                RobotState.motors_tor_des = RobotState.wbc_tauJointRes.flatten()

            pvtCtr.dataBusRead(RobotState)
            if (sim_time <= 3):
                pvtCtr.calMotorsPVT(100.0/1000.0/180.0*3.1415)
            else:
                pvtCtr.setJointPD(100,10,"J_ankle_l_pitch")
                pvtCtr.setJointPD(100,10,"J_ankle_l_roll")
                pvtCtr.setJointPD(100,10,"J_ankle_r_pitch")
                pvtCtr.setJointPD(100,10,"J_ankle_r_roll")
                pvtCtr.setJointPD(1000,100,"J_knee_l_pitch")
                pvtCtr.setJointPD(1000,100,"J_knee_r_pitch")
                pvtCtr.calMotorsPVT()
            pvtCtr.dataBusWrite(RobotState)

            # mj_interface.setMotorsTorque(RobotState.motors_tor_out);

            logger.start_new_line()
            logger.record_item_data("simTime", sim_time)
            logger.record_item_data("motors_pos_cur",RobotState.motors_pos_cur)
            logger.record_item_data("motors_vel_cur",RobotState.motors_vel_cur)
            logger.record_item_data("rpy",RobotState.rpy)
            logger.record_item_data("fL",RobotState.fL)
            logger.record_item_data("fR",RobotState.fR)
            logger.record_item_data("basePos",RobotState.basePos)
            logger.record_item_data("baseLinVel",RobotState.baseLinVel)
            logger.record_item_data("baseAcc",RobotState.baseAcc)
            logger.record_item_data("baseAngVel",RobotState.baseAngVel)
            logger.finish_line()

            print("rpyVal=[%.5f, %.5f, %.5f]\n", RobotState.rpy[0], RobotState.rpy[1], RobotState.rpy[2])
            print("gps=[%.5f, %.5f, %.5f]\n", RobotState.basePos[0], RobotState.basePos[1], RobotState.basePos[2])
            print("vel=[%.5f, %.5f, %.5f]\n", RobotState.baseLinVel[0], RobotState.baseLinVel[1], RobotState.baseLinVel[2])
                
        if mj_data.time >= sim_end_time:
            break
        
        uiController.updateScene()
            
    uiController.close()

    return 0

    # //    // free visualization storage
    # glfw.terminate()
# if __name__ == "__main":
main()